# 用于绘制boxes json 中的道路标线 及其他元素
import json
import matplotlib.pyplot as plt
import numpy as np
from matplotlib.patches import Polygon
from collections import defaultdict
import matplotlib.cm as cm
import matplotlib.colors as colors



route_cnt = 0

def load_json(boxes_path):
    # measurements = json.load(open(measure_path))
    boxes = json.load(open(boxes_path))
    return boxes

def process_json(boxes):
    # 车辆存储: 类别, x, y, yaw, extent_x, exten_y, vel, id
    # route 存储: 类别, x, y , yaw, extent_x, extent_y, id
    # lane 是点, 存储: 类别, x, y,  yaw, lane_id, point_id,
    
    agents = []
    ego = []
    routes = []
    lanes = []
    
    ego_data = boxes[0]
    ego.append([0.0, ego_data['position'][0], ego_data['position'][1], ego_data['yaw'], 
                     ego_data['extent'][1], ego_data['extent'][0], ego_data['speed'], ego_data['id']])
    for label_data in boxes[1:]:
        if label_data['class'] == 'Car':
            agents.append([1.0, label_data['position'][0], label_data['position'][1], label_data['yaw'], 
                           label_data['extent'][1], label_data['extent'][0], label_data['speed'], label_data['id']])
        elif label_data['class'] == 'Route':
            routes.append([2.0, label_data['position'][0], label_data['position'][1], label_data['yaw'], 
                           label_data['extent'][1], label_data['extent'][0], label_data['id']])
        elif label_data['class'] == 'Lane':
            lanes.append([3.0, label_data['position'][0], label_data['position'][1], label_data['yaw'], 
                           label_data['lane_id'], label_data['point_id']])
    return ego, agents, routes, lanes


def calc_rect_points(x, y, yaw, x_extent, y_extent):
     # 矩形的四个顶点（相对于中心点）
    vertices = np.array([
        [x_extent, y_extent],   # 右上角
        [-x_extent, y_extent],  # 左上角
        [-x_extent, -y_extent], # 左下角
        [x_extent, -y_extent]   # 右下角
    ])
    
    # 旋转矩阵
    rotation_matrix = np.array([
        [np.cos(yaw), -np.sin(yaw)],
        [np.sin(yaw), np.cos(yaw)]
    ])
    
    # 应用旋转
    rotated_vertices = np.dot(vertices, rotation_matrix.T)
    
    # 平移到中心点
    rotated_vertices += np.array([x, y])
    return rotated_vertices


def collect_lane(labels):
    # 按 lane_id 分组
    lane_groups = defaultdict(list)
    for point in labels:
        lane_groups[point[-2]].append(point)

    # 对每个 lane_id 内的点按 point_id 排序
    for lane_id in lane_groups:
        lane_groups[lane_id] = sorted(lane_groups[lane_id], key=lambda p: p[-1])

    return lane_groups
    

def draw_box(ax, label):
    if label[0] == 0.0:
        points = calc_rect_points(label[1], label[2], label[3], label[4], label[5])
        # 创建并添加多边形
        rect_patch = Polygon(points, closed=True, 
                         fill=True, color='red')
        ax.add_patch(rect_patch)
    
    elif label[0] == 1.0:
        points = calc_rect_points(label[1], label[2], label[3], label[4], label[5])
        # 创建并添加多边形
        rect_patch = Polygon(points, closed=True, 
                         fill=True, color='green')
        ax.add_patch(rect_patch)
    
    elif label[0] == 2.0:
        points = calc_rect_points(label[1], label[2], label[3], label[4], label[5])
        # 创建并添加多边形
        global route_cnt
        color_route = (10.0*route_cnt/255, 0.0, 1.0)
        rect_patch = Polygon(points, closed=True, 
                         fill=True, color=color_route, alpha=0.5)
        ax.add_patch(rect_patch)
        if route_cnt < 20:
            route_cnt += 1
    print("add patch")


def draw_lines(ax, lanes):
    cnt = 0
    id_to_cnt = dict()
    for lane_id, points in lanes.items():
        
        # if lane_id not in [65]:
        #     continue
        
        if len(points) <= 4 :
            continue
        id_to_cnt[lane_id] = cnt
        cnt += 1
    
    # 创建一个颜色映射
    norm = colors.Normalize(vmin=0, vmax=cnt)  # 假设最多100个不同的lane_id
    cmap = cm.get_cmap('viridis')  # 可以选择其他颜色映射，如 'jet', 'rainbow' 等  
    for lane_id, points in lanes.items():
        if lane_id not in id_to_cnt:
            continue
        print(f"lane id is {lane_id}, length is {len(points)}")
        color = cmap(norm(id_to_cnt[lane_id]))  # 直接从颜色映射获取颜色
        # color = (0.0, min(cnt/25.5, 1.0), 1.0)
        # cnt += 1
        x = [p[1] for p in points]
        y = [p[2] for p in points]
        ax.plot(x, y, '-', color=color, label=f"lane {lane_id}")
    

if __name__ == "__main__":
    boxes = load_json('/home/li/code/PlanT/data/PlanT_data_1/coke_s3_dataset_12_08/Routes_Town01_Scenario3_Seed2001/Town01_Scenario3_route0_09_27_17_24_33/boxes/0000.json')
    ego, agents, routes, lanes = process_json(boxes)
    lanes_group = collect_lane(lanes)
    
    
    # 在函数外部创建Figure和Axes
    fig, ax = plt.subplots()
    for part in [ego, agents, routes]:
        for data in part:
            draw_box(ax, data)
    draw_lines(ax, lanes_group)
    
    ax.set_xlabel('X')
    ax.set_ylabel('Y')
    plt.legend()
    ax.set_xlim(-30, 30)
    ax.set_ylim(-30, 30)
    ax.set_aspect('equal')
    plt.show()
    
    
    